/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.cdiggins.com
*/

#ifndef CVML_OPTIMIZER_HPP
#define CVML_OPTIMIZER_HPP

// Can only be called before compaction and tail-calls.
void CountSubRefs(ByteCode* bc)
{
	// reset the reference count of each subroutine
    for (int i=0; i < bc->GetNumSubs(); ++i) {
		bc->GetSub(i)->callcnt = 0;
		bc->GetSub(i)->pushcnt = 0;
	}

	// Count all references to each subroutine
    for (int i=0; i < bc->GetNumSubs(); ++i) {
        Subroutine* sub = bc->GetSub(i);
        for (Instruction* cur = sub->first; cur != NULL; cur = cur->next) {
            if (cur->opcode == op_push_sub) {
				Subroutine* sub2 = cur->GetSub();
				if (cur->next != NULL && cur->next->opcode == op_apply) 
					sub2->callcnt++;
				else
	                sub2->pushcnt++;
			}
        }
    }

	/*
	printf("Subroutine call counts\n");
    for (int i=0; i < bc->GetNumSubs(); ++i) {
		Subroutine* sub = bc->GetSub(i);
		fprintf(OUTPUT, "%s, push count %d\n", sub->name.c_str(), sub->pushcnt);
		fprintf(OUTPUT, "%s, call count %d\n", sub->name.c_str(), sub->callcnt);
	}
	printf("\n");
	*/
}

bool RemoveUnusedSubs(ByteCode* bc) 
{
	CountSubRefs(bc);

	bool r = false;
    for (int i=bc->GetNumSubs() - 1; i >=0; --i) {
        // We don't remove user defined functions. They aren't called
		// from within but the expectation is that they are exported. 
        Subroutine* sub = bc->GetSub(i);
		if (!bc->IsExported(sub) && sub->RefCnt() == 0) {
            bc->RemoveSub(i);
			r = true;
        }
    }

	return r;
}

// The "visited" parameter is important to assure we don't get caught chasing down 
// for a reference within a recursive function/
bool ReferencesSub(ByteCode* bc, Subroutine* sub, std::vector<Subroutine*>& visited) {
    for (Instruction* i = sub->first; i != NULL; i = i->next) {
        if (i->IsSubRefTo(sub)) {
            return true;
        }
        if (i->IsSubRef()) {
            Subroutine* tmp = i->GetSub();
            if (!Contains(visited, tmp)) {
                visited.push_back(tmp);
                if (ReferencesSub(bc, tmp, visited)) {
                    visited.pop_back();
                    return true;
                }
                visited.pop_back();
            }
        }
    }
    return false;
}

bool IsRecursive(ByteCode* bc, Subroutine* sub) {
    std::vector<Subroutine*> visited;
    return ReferencesSub(bc, sub, visited);
}

Instruction* CopySubroutineBody(Subroutine* sub) 
{
	Instruction* src = sub->first;
	Instruction* copy = new Instruction(*src);
	while (src->next != NULL) {
		copy = copy->AddInstruction(*src);
		src = src->next;
	}
	Assert(src->opcode == op_ret);
	return copy;
}

void InlineSub(ByteCode* bc, Subroutine* sub) {
	for (int i=0; i < bc->GetNumSubs(); ++i) {
		Subroutine* sub2 = bc->GetSub(i);
		// Prevent self recursive inlining.
		if (sub2 != sub) {
			for (Instruction* tmp = sub2->first; tmp != NULL; tmp = tmp->next) {
				if (tmp->IsCallTo(sub)) {
					tmp->SetNoop();
					tmp->next->SetNoop();
					Instruction* src = sub->first;
					Instruction* copy = CopySubroutineBody(sub);
					copy->next = tmp->next;
					tmp->next = copy;
				} 
			}
		}
	}
}

int SubroutineSize(Subroutine* sub) {
	int n = 0;
	for (Instruction* i = sub->first; i != NULL; i = i->next) {
		if (i->opcode != op_noop && i->opcode != op_ret)
			++n;
	}
	return n;
}

bool IsInlinable(ByteCode* bc, Subroutine* sub) {
	if (bc->IsExported(sub))
		return false;

	if (sub->callcnt == 1 && sub->pushcnt == 0)
		return true;

	// Check size of subroutine.
	// The heuristic varies depending on the size of the 
	// subroutine.
	if (sub->callcnt == 2) {
		if (SubroutineSize(sub) <= 3)
			return true;
	}
	else if (sub->callcnt > 2) {
		if (SubroutineSize(sub) <= 2) 
			return true;
	}

	return false;
}

bool InlineSubs(ByteCode* bc) {
	CountSubRefs(bc);
	bool r = false;
    for (int i=bc->GetNumSubs() - 1; i >=0; --i) {
        // We don't remove user defined functions
        Subroutine* sub = bc->GetSub(i);
		if (IsInlinable(bc, sub)) {
			InlineSub(bc, sub);
			bc->RemoveSub(i);
			r = true;
		}
    }
	return r;
}

void ReplaceCallsWith(ByteCode* bc, Subroutine* from, Subroutine* to) {
    for (int i=0; i < bc->GetNumSubs(); ++i) {
        Subroutine* sub = bc->GetSub(i);
        for (Instruction* cur = sub->first; cur != NULL; cur = cur->next)  {
            if (cur->IsSubRef()) {
                if (cur->GetSub() == from) {
                    cur->operand = reinterpret_cast<int>(to);
                }
            }
        }
    }
}

bool DoInstrMatch(Instruction* i, Instruction* j) {
	return i->Matches(*j);
}

bool DoSeqMatch(InstrSeq seq1, InstrSeq seq2) {
	if (seq1.Empty()) {
		if (seq2.Empty())
			return true;
		else
			return false;
	}
	if (DoInstrMatch(seq1.begin, seq2.begin)) {
		if (DoSeqMatch(seq1.Tail(), seq2.Tail())) 
			return true;
		else
			return false;
	}
	return false;
}

bool DoSubsMatch(Subroutine* sub1, Subroutine* sub2) {
	return DoSeqMatch(InstrSeq(sub1->first, NULL), InstrSeq(sub2->first, NULL));
}

void ReplaceDuplicateSubs(ByteCode* bc) {
    for (int i=0; i < bc->GetNumSubs() - 1; ++i) {
		Subroutine* sub1 = bc->GetSub(i);
		if (!bc->IsExported(sub1)) {
			for (int j=i+1; j < bc->GetNumSubs(); ++j) {
				Subroutine* sub2 = bc->GetSub(j);
				if (!bc->IsExported(sub2) && DoSubsMatch(sub1, sub2)) {
					ReplaceCallsWith(bc, sub1, sub2);
				}
			}
		}
    }
}

/*
void ExtendSubPrefixes(ByteCode* bc) {
    for (int i=0; i < bc->GetNumSubs(); ++i) {
        std::vector<Instruction*> calls;
        Subroutine* sub = bc->GetSub(i);
        IdentifyCalls(bc, &calls, sub);
        
        if (calls.size() > 0 && calls[0]->prev != NULL) {
            Instruction* prefix = calls[0]->prev;
            bool bmatch = true;
        
            for (size_t j=1; j < calls.size(); ++j) {
                Instruction* prefix2 = calls[1]->prev;                    
                if (prefix2 == NULL || !prefix2->Equals(*prefix))
                    bmatch = false;
            }
            
            if (bmatch) {
                sub->first = sub->first->AddInstructionBefore(*prefix);
                
                // Important: do this after!
                for (size_t j=0; j < calls.size(); ++j) {
                    calls[j]->prev->SetNoop();
                }
            }                    
        }
    }
}

void ExtendSubSuffixes(ByteCode* bc) {
    for (int i=0; i < bc->GetNumSubs(); ++i) {
        std::vector<Instruction*> calls;
        Subroutine* sub = bc->GetSub(i);
        IdentifyCalls(bc, &calls, sub);
        
        if (calls.size() > 0 && calls[0]->next != NULL) {
            Instruction* suffix = calls[0]->next;
            bool bmatch = true;
        
            for (size_t j=1; j < calls.size(); ++j) {
                Instruction* suffix2 = calls[1]->next;                    
                if (suffix2 == NULL || !suffix2->Equals(*suffix))
                    bmatch = false;
            }
            
            if (bmatch) {
                Instruction* last = sub->GetLastInstr();
                last->next = last->AddInstruction(*suffix);
                
                // Important: do this after!
                for (size_t j=0; j < calls.size(); ++j) {
                    calls[j]->next->SetNoop();
                }
            }                    
        }
    }
}
*/

// After the call to this function, all subroutines are referred to by an index and not 
// their memory address 
void ResolveSubroutines(ByteCode* bc)
{
	CountSubRefs(bc);

	// sort the sub by count (bubblesort)
	for (int i=0; i < bc->GetNumSubs(); ++i) {
		for (int j=i+1; j < bc->GetNumSubs(); ++j) {
			if (bc->GetSub(j)->RefCnt() > bc->GetSub(i)->RefCnt()) {
				bc->SwapSubs(i, j);
			}
		}
	}

	// Make sure the sort works
	for (int i=1; i < bc->GetNumSubs(); ++i)
		Assert(bc->GetSub(i)->RefCnt() <= bc->GetSub(i-1)->RefCnt());

    // change sub references to indexes
    for (int i=0; i < bc->GetNumSubs(); ++i) {
        Subroutine* sub = bc->GetSub(i);
        for (Instruction* cur = sub->first; cur != NULL; cur = cur->next)  {            
            // NOTE: This won't work if we have already performed peephole 
            // optimization. We have to be using the O1 instruction set.
            if (cur->IsSubRef()) {
				Subroutine* tmp = cur->GetSub();
                int id = bc->GetIndexOf(tmp);
                cur->operand = id;
            }
        }
    }
}

/*
void Optimize(ByteCode* bc)
{   
	bc->StripNoops();
	bc->WithEachSub(OptimizeZaps);
    PrintByteCodeAndStats("Optimized zap operations", bc);

	// Merge similar subroutines 
	//ProceduralAbstraction(bc);
    //PrintByteCodeAndStats("Performed parameterized procedural abstraction", bc);

	while (RemoveUnusedSubs(bc))
    { }
    PrintByteCodeAndStats("Removed unused functions", bc);

	bc->StripNoops(); 
	ReplaceDuplicateSubs(bc);
	while (RemoveUnusedSubs(bc))
    { }
    PrintByteCodeAndStats("Replaced duplicate subroutines", bc);

	bc->StripNoops(); 
	GenerateSubs(bc);
    PrintByteCodeAndStats("New subroutines created", bc);

	InlineSingleUseSubs(bc);
    PrintByteCodeAndStats("Inlined single-use functions", bc);

	Peephole(bc);
	PrintByteCodeAndStats("Peephole optimization applied", bc);

	MakeTailCalls(bc);
	PrintByteCodeAndStats("Performed tail call optimizations", bc);

	// One of the last things we do, is replace subroutine pointers
	// with indexes. In so doing we reorder the subroutines by 
	// frequency of call
    ResolveSubroutines(bc);
    PrintByteCodeAndStats("Subroutines reordered by frequency of call", bc);
    
    //CompactOps(bc);
    //PrintByteCodeAndStats("Final pass with peephole optimizer", bc);

	Peephole(bc);
	PrintByteCodeAndStats("Peephole optimization applied", bc);

	CompactInstructions(bc);
	PrintByteCodeAndStats("Constant folding", bc);
}
*/

void MakeTailCalls(ByteCode* bc) 
{
	for (int i=0; i < bc->GetNumSubs(); ++i) {
		for (Instruction* j = bc->GetSub(i)->first; j != NULL && j->next != NULL; j = j->next) {
			if (j->next->opcode == op_ret) 
			{
				switch (j->opcode)
				{
				case op_apply: j->opcode = op_tail_apply; j->next->SetNoop(); break;
				case op_if_else: j->opcode = op_if_else; j->next->SetNoop(); break;
				case op_if: j->opcode = op_tail_if; j->next->SetNoop(); break;
				//?? case op_dip: j->opcode = op_tail_dip; j->next->SetNoop(); break;
				//?? case op_sap: j->opcode = op_tail_sap; j->next->SetNoop(); break;
				}
			}
		}
	}
	bc->StripNoops();
}


#endif
